import datetime
import time
import pandas as pd
import tushare as ts
from tqdm import tqdm
from tools.utils import load_config
from tools.utils import loc_collection


def ini_all_future_contracts(cfg, contract_type="普通合约"):
    """
    从tushare下载所有合约数据保存至mongodb ts_code_future数据库中
    :return:
    """
    today = str(datetime.date.today())
    print(f"contract_type{contract_type}")
    contract_type_map = {"普通合约": 1, "主力连续合约": 2}
    pro = ts.pro_api(cfg['tushare'])
    for exchange, e_dic in cfg['exchange'].items():
        exchange_ts = cfg['exchange_map'][exchange]
        exchange_df = pro.fut_basic(exchange=exchange_ts, fut_type=contract_type_map[contract_type])
        exchange_df.index = exchange_df['ts_code']
        for fut_code, ts_df in tqdm(exchange_df.groupby(['fut_code']), ncols=80):
            fut_collection = loc_collection('tushare_future', fut_code)
            for ts_code in ts_df['ts_code'].unique():
                existed = list(fut_collection.find({'ts_code': ts_code}))
                if len(existed) > 0:
                    print("exists, inserted before")
                else:
                    daily_data = pro.fut_daily(ts_code=ts_code)
                    if len(daily_data) > 0:
                        daily_data = daily_data.sort_values(['trade_date'])
                        dic_data = {k: v.to_list() for k, v in daily_data.items()}
                        update = False
                        if contract_type == '主力连续合约':
                            update = True
                        else:
                            delist_date = exchange_df.loc[ts_code, 'delist_date']
                            if str(pd.to_datetime(str(delist_date)))[:10] > today:
                                update = True
                        insert_data = {
                            'ts_code': dic_data['ts_code'][0],
                            't': [str(pd.to_datetime(str(s)))[:10] for s in dic_data['trade_date']],
                            'start_t': dic_data['trade_date'][0],
                            'end_t': dic_data['trade_date'][-1],
                            'update': update,
                        }
                        del dic_data['ts_code']
                        for k, v in dic_data.items():
                            insert_data[k] = v

                        if len(insert_data) > 0:
                            fut_collection.insert_one(insert_data)
                            print('inserting data', ts_code, daily_data.shape)
                            time.sleep(0.5)


if __name__ == "__main__":
    # ini运行一次就够了
    conf = load_config()
    conf['exchange'] = {'CZC': {'甲醇': ['MA.CZC', 'ma_zce']}}
    ini_all_future_contracts(conf, "主力连续合约")
